import numpy as np 
import tensorflow as tf 
import cv2 as cv 
import os
import sklearn
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt 
from keras.layers import Input,Conv2D,BatchNormalization,Activation,MaxPool2D,merge,Flatten,Dense,concatenate,Add
from keras import optimizers
from keras.callbacks import ModelCheckpoint
import math
from keras.models import Model,Sequential
from keras.preprocessing.image import img_to_array,load_img
from PIL import Image


def conv_block(input_tensor, filters):
    filter1, filter2, filter3 = filters

    x = Conv2D(filter1,(1,1),strides=1)(input_tensor)
    x = BatchNormalization(axis=-1)(x)
    x = Activation('tanh')(x)

    x = Conv2D(filter2,(3,3),strides=1,padding='same')(x)
    x = BatchNormalization(axis=-1)(x)
    x = Activation('selu')(x)

    x = Conv2D(filter3,(1,1),strides=1)(x)
    x = BatchNormalization(axis=-1)(x)
    x = Activation('selu')(x)

    y = Conv2D(filter3,(1,1),strides=1)(input_tensor)
    y = BatchNormalization(axis=-1)(y)
    y = Activation('elu')(y)

    # out = merge([x,y],mode='sum')
    #out=concatenate([x,y],axis=3)
    out=Add()([x,y])
    z = Activation('tanh')(out)

    return z

def identity_block(input_tensor, filters):


    filter1, filter2, filter3 = filters

    x = Conv2D(filter1,(1,1),strides=1)(input_tensor)
    x = BatchNormalization(axis=-1)(x)
    x = Activation('tanh')(x)

    x = Conv2D(filter2,(3,3),strides=1,padding='same')(x)
    x = BatchNormalization(axis=-1)(x)
    x = Activation('selu')(x)

    x = Conv2D(filter3,(1,1),strides=1)(x)
    x = BatchNormalization(axis=-1)(x)
    x = Activation('selu')(x)

    # y = Conv2D(filter3,(1,1),strides=1)(input_tensor)
    # y = BatchNormalization(axis=-1)(y)
    # y = Activation('elu')(y)

    # out = merge([x,input_tensor],mode='sum')
    #out=concatenate([x,y],axis=3)
    out=Add()([x,input_tensor])
    z = Activation('tanh')(out)
    return z


def resnet_model(out_class, input_shape):

    inputs = Input(shape=input_shape) #1,3,224,224

    #
    x = Conv2D(64, (7, 7), strides=2, padding='same',data_format="channels_first")(inputs) #conv1  1,64,112,112
    
    x = BatchNormalization(axis=-1)(x) #bn_conv1
    x = Activation('tanh')(x) #conv1_relu

    x = MaxPool2D(pool_size=(3,3),strides=2)(x) # 1,64,56,56

    # block1  (64,64,256) 1,2 in:1,64,56,56
    x = conv_block(x, [64, 64, 256]) #out=1,256,56,56
    x = identity_block(x, [64, 64, 256]) #out=1,256,56,56
    x = identity_block(x, [64, 64, 256]) #out=1,256,56,56

    # block2  (128,128,512) 1,3 in:1,256,56,56
    x = conv_block(x, [128,128,512]) #out=1,512,28,28
    x = identity_block(x, [128,128,512]) #out=1,512,28,28
    x = identity_block(x, [128,128,512]) #out=1,512,28,28
    x = identity_block(x, [128, 128, 512])  # out=1,512,28,28

    # block 3 (256,256,1024) 1,5 in:1,512,28,28
    x = conv_block(x, [256,256,1024])  # out=1,1024,14,14
    x = identity_block(x, [256, 256, 1024])  # out=1,1024,14,14
    x = identity_block(x, [256, 256, 1024])  # out=1,1024,14,14
    x = identity_block(x, [256, 256, 1024])  # out=1,1024,14,14
    x = identity_block(x, [256, 256, 1024])  # out=1,1024,14,14
    x = identity_block(x, [256, 256, 1024])  # out=1,1024,14,14

    # block 4 (512,512,2048) 1,2 in:1,1024,14,14
    x = conv_block(x, [512,512,2048])  # out=1,2048,7,7
    x = identity_block(x, [512, 512, 2048])  # out=1,2048,7,7
    x = identity_block(x, [512, 512, 2048])  # out=1,2048,7,7

    # maxpool kernel_size=7, stride=1 out=1,2048,1,1
    x = MaxPool2D(pool_size=(7, 7), strides=1)(x)

    # flatten
    x = Flatten()(x)

    # # Dense
    # x = Dense(1000)(x) # out=1,1000

    # Dense,这里改造了一下，适应cifar10
    x = Dense(out_class)(x)  # out=1,1000

    # out = Activation('softmax')(x)
    # out=Activation('sigmoid')(x)
    model = Model(inputs=inputs, outputs=x)
    model.summary()
    adam=optimizers.Adam(lr=0.01)

    model.compile(optimizer="sgd",loss="mae",metrics=["accuracy"])
    return model